import requests
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
import re

# 读取文件并分行处理
with open("document.txt", "r", encoding="utf-8") as file:
    lines = file.readlines()
texts = [re.sub(r'\s+', ' ', line).strip() for line in lines if line.strip()]  # 清理每行数据

# 通过 Xinference 生成向量
model_url = "http://localhost:9997/v1/models"
payload = {"model_name": "bge-small-en-v1.5", "model_type": "embedding"}
response = requests.post(model_url, json=payload)
print("Model load response:", response.status_code, response.text)
model_info = response.json()
if "model_uid" not in model_info:
    raise ValueError(f"Failed to load model: {response.text}")
model_uid = model_info["model_uid"]

embed_url = "http://localhost:9997/v1/embeddings"
embeddings = []
for text in texts:
    payload = {"model": model_uid, "input": text}
    embedding = requests.post(embed_url, json=payload).json()["data"][0]["embedding"]
    embeddings.append(embedding)

# 连接 Milvus 并存储数据
connections.connect(host='localhost', port='19530')
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=384),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535)
]

# 如果集合已存在，先删除
try:
    collection = Collection(name="text_collection")
    collection.drop()
except:
    pass

# 创建新集合并插入数据
collection = Collection(name="text_collection", schema=CollectionSchema(fields=fields))
ids = list(range(1, len(texts) + 1))  # 生成 ID
collection.insert([ids, embeddings, texts])
collection.create_index("embedding", {"metric_type": "COSINE", "index_type": "IVF_FLAT", "params": {"nlist": 1024}})
collection.load()

print("Inserted", collection.num_entities, "entities")